#!/usr/bin/env python3
"""
SDXL LoRA训练节点 - 基于基类实现
"""

# 导入ComfyUI模块
try:
    from server import PromptServer
except ImportError:
    PromptServer = None

# 导入基类
from .base_lora_trainer_node import BaseLoRATrainerNode

# 导入常量
from .lora_trainer_utils.constants import AUX_ID,SDXL_LORA_TRAINER_CATEGORY

class SDXLLoRATrainer(BaseLoRATrainerNode):
    """SDXL LoRA训练器节点"""
    
    def __init__(self):
        super().__init__("train_network", "NetworkTrainer")
    
    def _send_progress_event(self, progress_data: dict):
        """SDXL训练器自定义进度事件发送
        
        Args:
            progress_data (dict): 进度数据字典
        """
        try:
            PromptServer.instance.send_sync("sdxl_training_progress", progress_data)
            print(f"[DEBUG] ✅ SDXL进度事件已发送: {progress_data}")
        except Exception as e:
            print(f"[WARNING] 发送SDXL进度更新失败: {e}")
    
    @classmethod
    def INPUT_TYPES(cls):
        """定义输入类型"""
        return {
            "required": {
                "training_params": ("SDXL_TRAINING_PARAMS", {
                    "tooltip": "SDXL训练参数配置"
                }),
            },
            "hidden": {
                "node_id": "UNIQUE_ID"
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING", "STRING")
    RETURN_NAMES = ("训练状态", "输出路径", "训练日志")
    FUNCTION = "train_sdxl_lora"
    CATEGORY = SDXL_LORA_TRAINER_CATEGORY
    OUTPUT_NODE = True

    def train_sdxl_lora(self, training_params, node_id=None):
        """训练SDXL LoRA模型"""
        return self.train_lora(training_params, node_id)

class SimpleSDXLLoRATrainer(BaseLoRATrainerNode):
    """简单SDXL LoRA训练器节点"""
    
    def __init__(self):
        super().__init__("train_network", "NetworkTrainer")
    
    def _send_progress_event(self, progress_data: dict):
        """SDXL训练器自定义进度事件发送
        
        Args:
            progress_data (dict): 进度数据字典
        """
        try:
            PromptServer.instance.send_sync("sdxl_training_progress", progress_data)
            print(f"[DEBUG] ✅ SDXL进度事件已发送: {progress_data}")
        except Exception as e:
            print(f"[WARNING] 发送SDXL进度更新失败: {e}")
    
    @classmethod
    def INPUT_TYPES(cls):
        """定义输入类型"""
        return {
            "required": {
                "training_params": ("SDXL_TRAINING_PARAMS", {
                    "tooltip": "SDXL训练参数配置"
                }),
            },
            "hidden": {
                "node_id": "UNIQUE_ID"
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING", "STRING")
    RETURN_NAMES = ("训练状态", "输出路径", "训练日志")
    FUNCTION = "train_simple_sdxl_lora"
    CATEGORY = SDXL_LORA_TRAINER_CATEGORY
    OUTPUT_NODE = True

    def train_simple_sdxl_lora(self, training_params, node_id=None):
        """简单SDXL LoRA训练"""
        return self.train_lora(training_params, node_id)

# 导出节点
NODE_CLASS_MAPPINGS = {
    "SDXLLoRATrainer": SDXLLoRATrainer,
    "SimpleSDXLLoRATrainer": SimpleSDXLLoRATrainer
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SDXLLoRATrainer": "🔵 SDXL LoRA训练器",
    "SimpleSDXLLoRATrainer": "🔵 SDXL简单训练器"
} 